#include "nefunction.h"

extern ULONG_PTR ImFindNoExportedFunction(
    IN unsigned char* lpstrMk1,
    IN ULONG Mk1Size,
    IN unsigned char* lpstrMk2,
    IN ULONG Mk2Size,
    IN ULONG Mk1ToMk2,
    IN ULONG Mk1ToHead
);

NTSTATUS PsSuspendThread(
    IN PETHREAD Thread,
    OUT PULONG PreviousSuspendCount OPTIONAL
)
{
    typedef NTSTATUS(NTAPI* FnPsSuspendThread)(
        IN PETHREAD Thread,
        OUT PULONG PreviousSuspendCount OPTIONAL
        );

    static FnPsSuspendThread lpFnPsSuspendThread = NULL;

    if (!lpFnPsSuspendThread)
    {
        //  WIN10 x64
        //UCHAR s1[] = { 0x48, 0x89, 0x54, 0x24, 0x10, 0x48, 0x89, 0x4c, 0x24, 0x08 };
        //UCHAR s2[] = { 0x66, 0xff, 0x8e, 0xe4, 0x01, 0x00, 0x00, 0x4c, 0x8d, 0xb9, 0xc8, 0x06, 0x00, 0x00 };

        //  WIN7 x64
        UCHAR s1[] = { 0x48, 0x89, 0x54, 0x24, 0x10, 0x48, 0x89, 0x4c, 0x24, 0x08 };
        UCHAR s2[] = { 0x80, 0x00, 0x00, 0x00, 0x66, 0x41 };
        lpFnPsSuspendThread = (FnPsSuspendThread)ImFindNoExportedFunction(s1, sizeof(s1), s2, sizeof(s2), 0x2E, 0);
    }

    if (lpFnPsSuspendThread)
        return lpFnPsSuspendThread(Thread, PreviousSuspendCount);

    return STATUS_UNSUCCESSFUL;
}

NTSTATUS PsResumeThread(
    IN PETHREAD Thread,
    OUT PULONG PreviousSuspendCount OPTIONAL
)
{
    typedef NTSTATUS(NTAPI* FnPsResumeThread)(
        IN PETHREAD Thread,
        OUT PULONG PreviousSuspendCount OPTIONAL
        );

    static FnPsResumeThread lpFnPsResumeThread = NULL;

    if (!lpFnPsResumeThread)
    {
        //  WIN10 x64
        //UCHAR s1[] = { 0x48, 0x89, 0x5c, 0x24, 0x08, 0x48, 0x89, 0x74, 0x24, 0x10 };
        //UCHAR s2[] = { 0x4c, 0x8b, 0x87, 0x20, 0x02, 0x00, 0x00, 0xb8, 0x00, 0x80, 0x00, 0x00, 0x41, 0x8b, 0x88, 0xfc, 0x06, 0x00, 0x00 };
        //lpFnPsResumeThread = (FnPsResumeThread)ImFindNoExportedFunction(s1, sizeof(s1), s2, sizeof(s2), 0x2A, 0);

        //  WIN7 x64
        UCHAR s1[] = { 0x40, 0x53, 0x48, 0x83, 0xec, 0x20 };
        UCHAR s2[] = { 0x74, 0x02, 0x89, 0x03 };
        lpFnPsResumeThread = (FnPsResumeThread)ImFindNoExportedFunction(s1, sizeof(s1), s2, sizeof(s2), 0x11, 0);
    }

    if (lpFnPsResumeThread)
        return lpFnPsResumeThread(Thread, PreviousSuspendCount);

    return STATUS_UNSUCCESSFUL;
}

NTSTATUS ZwGetNextThread(
    IN HANDLE ProcessHandle,
    IN HANDLE ThreadHandle,
    IN ACCESS_MASK DesiredAccess,
    IN ULONG HandleAttributes,
    IN ULONG Flags,
    OUT PHANDLE NewThreadHandle
)
{
    typedef NTSTATUS(NTAPI* FnZwGetNextThread)(
        IN HANDLE ProcessHandle,
        IN HANDLE ThreadHandle,
        IN ACCESS_MASK DesiredAccess,
        IN ULONG HandleAttributes,
        IN ULONG Flags,
        OUT PHANDLE NewThreadHandle
        );

    static FnZwGetNextThread lpFnZwGetNextThread = NULL;

    if (!lpFnZwGetNextThread)
    {
        //  WIN10 x64
        //UNICODE_STRING usZwGetNextThread = RTL_CONSTANT_STRING(L"ZwGetNextThread");
        //lpFnZwGetNextThread = (FnZwGetNextThread)MmGetSystemRoutineAddress(&usZwGetNextThread);

        //  WIN7 x64
        UCHAR s1[] = { 0x48, 0x8b, 0xc4, 0xfa, 0x48, 0x83, 0xec, 0x10 };
        UCHAR s2[] = { 0x50, 0xb8, 0xcf, 0x00, 0x00, 0x00 };
        lpFnZwGetNextThread = (FnZwGetNextThread)ImFindNoExportedFunction(s1, sizeof(s1), s2, sizeof(s2), 0x13, 0);
    }

    if (lpFnZwGetNextThread)
        return lpFnZwGetNextThread(ProcessHandle, ThreadHandle, DesiredAccess, HandleAttributes, Flags, NewThreadHandle);

    return STATUS_UNSUCCESSFUL;
}